// main.go
package main

import (
	"bytes"
	"compress/gzip"
	"context"
	"encoding/json"
	"encoding/xml"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"path/filepath"
	"sync"
	"syscall"
	"time"
)

// 常量定义
const (
	MaxURLsPerSitemap = 50000 // 每个站点地图文件最多包含的 URL 数量
	MaxUncompressedMB = 50 * 1024 * 1024 // 每个站点地图文件最大未压缩大小（50MB）
	XMLNS             = "http://www.sitemaps.org/schemas/sitemap/0.9" // XML 命名空间
	timeFormatFull    = "20060102-150405" // 时间戳格式：YYYYMMDD-HHMMSS
	stateFileName     = "state.json" // 状态文件名，用于记录 URL 的最后修改时间
)

// UrlEntry 表示一个 URL 条目
// XML 标签定义了序列化时的字段名称
type UrlEntry struct {
	Loc        string `xml:"loc"`        // 必填：URL 地址
	LastMod    string `xml:"lastmod,omitempty"`  // 选填：最后修改时间，格式为 YYYY-MM-DD
	ChangeFreq string `xml:"changefreq,omitempty"` // 选填：更新频率（always, hourly, daily, weekly, monthly, yearly, never）
	Priority   string `xml:"priority,omitempty"`   // 选填：优先级（0.0 - 1.0）
}

// SitemapMeta 表示一个站点地图文件的元信息
type SitemapMeta struct {
	Loc     string // 文件位置（路径）
	LastMod string // 最后修改时间
}

// Config 配置结构体，包含程序运行所需的所有参数
type Config struct {
	OutputDir     string // 输出目录路径
	TotalURLs     int    // 总 URL 数量（演示用，实际应从数据库流式获取）
	WorkerCount   int    // 并发写入的工作协程数
	BaseURL       string // 网站基础 URL，用于生成完整链接和 ping 操作
	ChunkSize     int    // 每个增量文件的 URL 分块大小
	IndexFileName string // 主索引文件名
	Keep          int    // 在主索引中保留的站点地图条目数量
	PingEnabled   bool   // 是否启用搜索引擎 ping 功能
}

// StateMap 状态映射，记录每个 URL 的最后修改时间
// 用于对比检测 URL 是否发生变化
type StateMap map[string]string

// ----------------- 工具函数 -----------------

// ensureDir 确保指定目录存在，如果不存在则创建
// 参数 dir: 目录路径
// 返回值: 错误信息（如果创建失败）
func ensureDir(dir string) error {
	return os.MkdirAll(dir, 0o755)
}

// ----------------- 生成示例 URL（实际项目中应替换为真实数据源） -----------------

// generateURL 生成第 i 个示例 URL 条目
// 参数 i: URL 的序号
// 参数 base: 基础 URL
// 返回值: 一个填充好的 UrlEntry 结构体
func generateURL(i int, base string) UrlEntry {
	// 示例 lastmod: 当前时间减去 (i%7) 天
	lm := time.Now().Add(-time.Duration(i%7) * 24 * time.Hour).Format("2006-01-02")
	return UrlEntry{
		Loc:        fmt.Sprintf("%s/item/%d", base, i), // 生成唯一 URL
		LastMod:    lm,                                 // 设置最后修改时间
		ChangeFreq: "weekly",                          // 固定更新频率为 weekly
		Priority:   fmt.Sprintf("%.1f", 0.5+float64(i%5)/10.0), // 优先级在 0.5~0.9 之间循环
	}
}

// ----------------- 状态文件的加载与保存 -----------------

// loadState 从指定路径加载状态文件
// 如果文件不存在，则返回一个空的状态映射
// 参数 path: 状态文件路径
// 返回值: 状态映射和错误信息
func loadState(path string) (StateMap, error) {
	var st StateMap
	data, err := os.ReadFile(path)
	if err != nil {
		if os.IsNotExist(err) {
			return make(StateMap), nil // 文件不存在，返回空状态
		}
		return nil, err
	}
	if len(data) == 0 {
		return make(StateMap), nil // 文件为空，返回空状态
	}
	if err := json.Unmarshal(data, &st); err != nil {
		return nil, err
	}
	return st, nil
}

// saveState 将状态映射安全地保存到指定路径
// 使用临时文件进行原子写入，避免写入过程中文件损坏
// 参数 path: 目标文件路径
// 参数 st: 要保存的状态映射
// 返回值: 错误信息（如果保存失败）
func saveState(path string, st StateMap) error {
	b, err := json.MarshalIndent(st, "", "  ") // 格式化 JSON
	if err != nil {
		return err
	}
	tmp := path + ".tmp" // 临时文件名
	if err := os.WriteFile(tmp, b, 0o644); err != nil {
		return err
	}
	return os.Rename(tmp, path) // 原子性重命名
}

// ----------------- 带限制的 XML 写入（支持 GZIP 压缩和分片） -----------------

// writeSitemapGzipWithLimit 将 URL 流写入带 GZIP 压缩的站点地图文件，并根据数量或大小自动分片
// 参数 ctx: 上下文，用于控制协程生命周期
// 参数 baseFilename: 输出文件的基础名称
// 参数 urls: URL 条目通道，作为数据源
// 参数 maxURLs: 单个文件最大 URL 数量
// 参数 maxBytes: 单个文件最大未压缩字节数
// 返回值: 生成的文件元信息列表和错误信息
func writeSitemapGzipWithLimit(ctx context.Context, baseFilename string, urls <-chan UrlEntry, maxURLs int, maxBytes int64) ([]SitemapMeta, error) {
	const xmlHeader = `<?xml version="1.0" encoding="UTF-8"?>` // XML 文件头
	var metas []SitemapMeta      // 存储生成的每个分片文件的元信息
	part := 1                    // 当前分片序号
	var f *os.File               // 当前输出文件句柄
	var gw *gzip.Writer          // GZIP 压缩写入器
	var enc *xml.Encoder         // XML 编码器
	var currentSize int64        // 当前文件已写入的近似大小（未压缩）
	var urlCount int             // 当前文件已写入的 URL 数量
	now := time.Now().UTC()      // 当前 UTC 时间，用于元信息

	// openNew 打开一个新的分片文件
	// 创建临时文件、GZIP 写入器和 XML 编码器，并写入 XML 头部和 <urlset> 开始标签
	// 返回值: 错误信息（如果打开失败）
	openNew := func() error {
		filename := fmt.Sprintf("%s-part%03d.xml.gz", baseFilename, part) // 构建文件名
		tmp := filename + ".tmp" // 临时文件名，用于原子写入
		var err error
		f, err = os.Create(tmp) // 创建临时文件
		if err != nil {
			return fmt.Errorf("create tmp %s: %w", tmp, err)
		}
		gw = gzip.NewWriter(f)     // 创建 GZIP 压缩写入器
		enc = xml.NewEncoder(gw)   // 创建 XML 编码器，输出到 GZIP 流
		currentSize = 0            // 重置当前文件大小计数
		urlCount = 0               // 重置当前文件 URL 计数
		// 注意：不再直接写入文件 f，而是通过 GZIP 流 gw 写入，确保整个内容被压缩
		_, _ = gw.Write([]byte(xmlHeader)) // 写入 XML 头部到压缩流
    currentSize += int64(len(xmlHeader))
		
		start := xml.StartElement{Name: xml.Name{Local: "urlset"}} // 构建 <urlset> 标签
		start.Attr = []xml.Attr{{Name: xml.Name{Local: "xmlns"}, Value: XMLNS}} // 添加 xmlns 属性
		_ = enc.EncodeToken(start) // 通过 XML 编码器写入开始标签
		_, _ = gw.Write([]byte("\n")) // 写入换行符
		currentSize += 50 // 近似估算已写入大小
		return nil
	}

	closeCurrent := func() error {
		if enc == nil {
			return nil
		}
		start := xml.StartElement{Name: xml.Name{Local: "urlset"}}
		_ = enc.EncodeToken(start.End())
		_ = enc.Flush()
		_ = gw.Close()
		_ = f.Close()
		filename := fmt.Sprintf("%s-part%03d.xml.gz", baseFilename, part)
		tmp := filename + ".tmp"
		final := filename
		if err := os.Rename(tmp, final); err != nil {
			return err
		}
		metas = append(metas, SitemapMeta{
			Loc:     filepath.ToSlash(final),
			LastMod: now.Format("2006-01-02"),
		})
		log.Printf("[INFO] wrote %s (urls=%d, size≈%d bytes)\n", final, urlCount, currentSize)
		part++
		return nil
	}

	if err := openNew(); err != nil {
		return nil, err
	}

	for {
		select {
		case <-ctx.Done():
			return metas, ctx.Err()
		case u, ok := <-urls:
			if !ok {
				// finish
				if err := closeCurrent(); err != nil {
					return metas, err
				}
				return metas, nil
			}
			// encode to buffer to measure uncompressed bytes
			var buf bytes.Buffer
			tmpEnc := xml.NewEncoder(&buf)
			_ = tmpEnc.Encode(u)
			_ = tmpEnc.Flush()
			b := buf.Bytes()

			if (currentSize + int64(len(b))) > maxBytes || urlCount >= maxURLs {
				if err := closeCurrent(); err != nil {
					return metas, err
				}
				if err := openNew(); err != nil {
					return metas, err
				}
			}

			if _, err := gw.Write(b); err != nil {
				return metas, err
			}
			if _, err := gw.Write([]byte("\n")); err != nil {
				return metas, err
			}
			currentSize += int64(len(b)) + 1
			urlCount++
		}
	}
}

// ----------------- main sitemap index append & cleanup -----------------

type sitemapIndexItem struct {
	Loc     string `xml:"loc"`
	LastMod string `xml:"lastmod,omitempty"`
}

type sitemapIndex struct {
	XMLName  xml.Name           `xml:"sitemapindex"`
	XMLNS    string             `xml:"xmlns,attr"`
	Sitemaps []sitemapIndexItem `xml:"sitemap"`
}

// append new entries and keep last `keep` items; returns total entries count after update.
func appendToSitemapIndex(indexPath string, newEntries []SitemapMeta, baseURL string, keep int) (int, error) {
	var idx sitemapIndex
	// Try read existing
	if data, err := os.ReadFile(indexPath); err == nil && len(data) > 0 {
		if err := xml.Unmarshal(data, &idx); err != nil {
			log.Printf("[WARN] failed to parse existing index: %v, will recreate", err)
			idx = sitemapIndex{}
		}
	}
	// append
	for _, e := range newEntries {
		item := sitemapIndexItem{
			Loc:     fmt.Sprintf("%s/%s", baseURL, filepath.Base(e.Loc)),
			LastMod: e.LastMod,
		}
		idx.Sitemaps = append(idx.Sitemaps, item)
	}
	// trim to keep last `keep`
	if keep > 0 && len(idx.Sitemaps) > keep {
		idx.Sitemaps = idx.Sitemaps[len(idx.Sitemaps)-keep:]
	}
	idx.XMLName = xml.Name{Local: "sitemapindex"}
	idx.XMLNS = XMLNS

	out, err := xml.MarshalIndent(idx, "", "  ")
	if err != nil {
		return 0, err
	}
	tmp := indexPath + ".tmp"
	if err := os.WriteFile(tmp, append([]byte(xml.Header), out...), 0o644); err != nil {
		return 0, err
	}
	if err := os.Rename(tmp, indexPath); err != nil {
		return 0, err
	}
	return len(idx.Sitemaps), nil
}

// delete files in dir that match patternPrefix and are NOT referenced in current index
func cleanupOldFiles(dir string, indexPath string, baseURL string) error {
	// read current index to get referenced basenames
	data, err := os.ReadFile(indexPath)
	if err != nil {
		return err
	}
	var idx sitemapIndex
	if err := xml.Unmarshal(data, &idx); err != nil {
		return err
	}
	keepSet := make(map[string]struct{})
	for _, s := range idx.Sitemaps {
		u, err := url.Parse(s.Loc)
		if err != nil {
			continue
		}
		keepSet[filepath.Base(u.Path)] = struct{}{}
	}
	// list files in dir that look like delta-*.xml.gz or sitemap-*.xml.gz or merged-*.xml.gz
	matches, _ := filepath.Glob(filepath.Join(dir, "*.xml.gz"))
	for _, f := range matches {
		b := filepath.Base(f)
		if _, ok := keepSet[b]; !ok {
			// delete
			if err := os.Remove(f); err != nil {
				log.Printf("[WARN] failed to remove %s: %v", f, err)
			} else {
				log.Printf("[INFO] removed old sitemap file: %s", f)
			}
		}
	}
	return nil
}

// ----------------- ping search engines -----------------

func pingSearchEngines(sitemapURL string) {
	pingEndpoints := []string{
		"https://www.google.com/ping?sitemap=" + url.QueryEscape(sitemapURL),
		"https://www.bing.com/ping?sitemap=" + url.QueryEscape(sitemapURL),
	}
	for _, endpoint := range pingEndpoints {
		log.Printf("[INFO] pinging: %s", endpoint)
		resp, err := http.Get(endpoint)
		if err != nil {
			log.Printf("[WARN] ping failed: %v", err)
			continue
		}
		body, _ := io.ReadAll(resp.Body)
		resp.Body.Close()
		log.Printf("[INFO] ping response: %d - %s", resp.StatusCode, string(bytes.TrimSpace(body)))
	}
}

// ----------------- diff logic -----------------

func diffURLs(old StateMap, all []UrlEntry) (changed []UrlEntry, newState StateMap) {
	newState = make(StateMap, len(all))
	for _, u := range all {
		newState[u.Loc] = u.LastMod
		if last, ok := old[u.Loc]; !ok || last != u.LastMod {
			changed = append(changed, u)
		}
	}
	return
}

// ----------------- orchestration -----------------

func runOnce(cfg Config) error {
	_ = ensureDir(cfg.OutputDir)
	statePath := filepath.Join(cfg.OutputDir, stateFileName)
	oldState, err := loadState(statePath)
	if err != nil {
		return fmt.Errorf("load state: %w", err)
	}

	// generate all URLs (replace this with DB streaming in prod)
	all := make([]UrlEntry, 0, cfg.TotalURLs)
	for i := 0; i < cfg.TotalURLs; i++ {
		all = append(all, generateURL(i, cfg.BaseURL))
	}

	changed, newState := diffURLs(oldState, all)
	if len(changed) == 0 {
		log.Println("[INFO] no URL changes detected, skipping sitemap generation")
		return nil
	}
	log.Printf("[INFO] detected %d new/updated URLs", len(changed))

	// prepare channels for writing (split changed into chunks and write with concurrency)
	// We'll create files named like: delta-YYYYMMDD-HHMMSS-<workerid>
	timestamp := time.Now().Format(timeFormatFull)
	// feed changed URLs into worker pipelines in chunks
	chunkSize := cfg.ChunkSize
	if chunkSize <= 0 || chunkSize > MaxURLsPerSitemap {
		chunkSize = MaxURLsPerSitemap
	}

	var metas []SitemapMeta
	var mu sync.Mutex
	var wg sync.WaitGroup
	sem := make(chan struct{}, cfg.WorkerCount)
	for i := 0; i < len(changed); i += chunkSize {
		end := i + chunkSize
		if end > len(changed) {
			end = len(changed)
		}
		wg.Add(1)
		sem <- struct{}{}
		go func(slice []UrlEntry, partIdx int) {
			defer wg.Done()
			defer func() { <-sem }()
			baseName := fmt.Sprintf("delta-%s-%03d", timestamp, partIdx)
			outPath := filepath.Join(cfg.OutputDir, baseName)
			// create url channel
			urlCh := make(chan UrlEntry, 256)
			go func() {
				defer close(urlCh)
				for _, u := range slice {
					// double-check dedupe in this chunk (simple map)
					urlCh <- u
				}
			}()
			// write
			pm, err := writeSitemapGzipWithLimit(context.Background(), outPath, urlCh, MaxURLsPerSitemap, MaxUncompressedMB)
			if err != nil {
				log.Printf("[ERROR] write sitemap chunk failed: %v", err)
				return
			}
			mu.Lock()
			metas = append(metas, pm...)
			mu.Unlock()
		}(changed[i:end], i/chunkSize+1)
	}
	wg.Wait()

	if len(metas) == 0 {
		log.Println("[ERROR] no metas produced")
		return nil
	}

	// append new metas to main index and keep only last cfg.Keep entries
	indexPath := filepath.Join(cfg.OutputDir, cfg.IndexFileName)
	totalEntries, err := appendToSitemapIndex(indexPath, metas, cfg.BaseURL, cfg.Keep)
	if err != nil {
		return fmt.Errorf("append index: %w", err)
	}
	log.Printf("[INFO] appended %d sitemap entries; main index now has %d entries", len(metas), totalEntries)

	// cleanup old files not referenced by index
	if err := cleanupOldFiles(cfg.OutputDir, indexPath, cfg.BaseURL); err != nil {
		log.Printf("[WARN] cleanup old files: %v", err)
	}

	// save new state
	if err := saveState(statePath, newState); err != nil {
		log.Printf("[WARN] save state failed: %v", err)
	} else {
		log.Printf("[INFO] state updated: %s", statePath)
	}

	// Ping search engines with index URL
	if cfg.PingEnabled {
		indexURL := fmt.Sprintf("%s/%s", cfg.BaseURL, cfg.IndexFileName)
		pingSearchEngines(indexURL)
	}

	return nil
}

// ----------------- main -----------------

func main() {
	var cfg Config
	flag.StringVar(&cfg.OutputDir, "out", "out", "output directory")
	flag.IntVar(&cfg.TotalURLs, "n", 50000, "total URL count (for demo; replace with DB)")
	flag.IntVar(&cfg.WorkerCount, "workers", 4, "concurrent workers for writing")
	flag.StringVar(&cfg.BaseURL, "base", "https://example.com", "site base URL")
	flag.IntVar(&cfg.ChunkSize, "chunk", MaxURLsPerSitemap, "chunk size per delta file (<=50000)")
	flag.StringVar(&cfg.IndexFileName, "index", "sitemap-index.xml", "main sitemap index filename")
	flag.IntVar(&cfg.Keep, "keep", 50, "how many sitemap entries to keep in main index")
	flag.BoolVar(&cfg.PingEnabled, "ping", true, "ping search engines after update")
	flag.Parse()

	_ = ensureDir(cfg.OutputDir)
	// simple signal handling for graceful shutdown while waiting (not strictly needed here)
	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
	defer stop()

	// single-run mode
	if err := runOnce(cfg); err != nil {
		log.Printf("[FATAL] run failed: %v", err)
		os.Exit(1)
	}

	// keep program alive briefly to allow signal cancellation if invoked interactively
	select {
	case <-ctx.Done():
	case <-time.After(500 * time.Millisecond):
	}
}
